from yacs.config import CfgNode

_cfg = CfgNode()
_cfg.desc = ""
_cfg.stage = 'train'  # train or eval or test
_cfg.device = 'cuda'  # cpu or cuda
_cfg.device_ids = ''
_cfg.output_dir = "outputs/veri776_b64_pvasff"
_cfg.seed = 42
_cfg.no_part_backward = False
_cfg.after_mask = False
_cfg.test_neck_feat = 'after'
_cfg.train_neck_feat = 'after'
_cfg.train = CfgNode()
_cfg.train.epochs = 120

_cfg.data = CfgNode()
_cfg.data.name = 'VeRi776'
_cfg.data.pkl_path = "preprocess/pkls/veri776.pkl"
_cfg.data.train_size = (224, 224)
_cfg.data.valid_size = (224, 224)
_cfg.data.pad = 10
_cfg.data.re_prob = 0.5
_cfg.data.with_mask = True
_cfg.data.test_ext = ''

_cfg.data.sampler = 'RandomIdentitySampler'
_cfg.data.batch_size = 2
_cfg.data.num_instances = 1

_cfg.data.train_num_workers = 8
_cfg.data.test_num_workers = 8

_cfg.model = CfgNode()
_cfg.model.ckpt_period = 10
_cfg.model.local_branches = 4
_cfg.model.pam = True
_cfg.model.attention = False

_cfg.backbone = CfgNode(new_allowed=True)
_cfg.backbone.name = 'resnet50'
_cfg.backbone.pretrained = True
_cfg.backbone.pretrained_path = ""
_cfg.backbone.suspend = -1
_cfg.backbone.apply_first_maxpool = False

# cbam parameters
_cfg.cbam = CfgNode()
_cfg.cbam.activation = 'LeakyReLU'
_cfg.cbam.params = ['negative_slope', 0.25]

# attention parameters
_cfg.attention = CfgNode()

'''
ExternalAttention SEAttention ChannelSpatialAttentionModule
'''
_cfg.attention.use_one = False
_cfg.attention.name = 'ChannelSpatialAttentionModule'
_cfg.attention.part_name = 'ExternalAttention'
# ----------ExternalAttention----------
_cfg.attention.ExternalAttention_spa_S = 2048
_cfg.attention.ExternalAttention_batch_S = 7 * 7
# ----------SEAttention----------
_cfg.attention.SEAttention_reduction = 16
# ----------ChannelSpatialAttentionModule----------
_cfg.attention.ChannelSpatialAttentionModule_reduction = 16
_cfg.attention.ChannelSpatialAttentionModule_kernel_size = 7
# ----------SKAttention----------
_cfg.attention.SKAttention_kernels = [1, 3, 5, 7]
_cfg.attention.SKAttention_reduction = 16
_cfg.attention.SKAttention_group = 1
_cfg.attention.SKAttention_L = 32
# ----------ECAAttention----------
_cfg.attention.ECAAttention_kernel_size = 3
# ----------DAModule----------
_cfg.attention.DAModule_kernel_size = 3

# pam parameters
_cfg.pam = CfgNode()
_cfg.pam.dropout_rate = .0
_cfg.pam.lowrank = False
_cfg.pam.rank = 1
_cfg.pam.fusion = True

_cfg.optim = CfgNode()
_cfg.optim.name = 'AdamW'  # SGD  AdamW
_cfg.optim.optim_kwargs = []
# base_lr and base_weight_decay must greater than zero
_cfg.optim.base_lr = 0.00035
_cfg.optim.base_weight_decay = 0.00005
# module lr or weight_decay equal to zero means it will follow base_lr or base_weight_decay
_cfg.optim.asff_lr = 0.
_cfg.optim.asff_weight_decay = 0.
_cfg.optim.backbone_lr = 0.
_cfg.optim.backbone_weight_decay = 0.
_cfg.optim.pam_lr = 0.
_cfg.optim.pam_weight_decay = 0.
_cfg.optim.bias_lr_factor = 1.
_cfg.optim.bias_decay_factor = 1.

_cfg.loss = CfgNode()
_cfg.loss.losses = ['id', 'supcon']  # supcon

# triplet loss parameters
_cfg.loss.triplet_margin = 0.3
_cfg.loss.normalize_feature = True

# id loss parameters
_cfg.loss.id_epsilon = 0.1

# parsing id loss parameters
_cfg.loss.pid_gamma = 1.
_cfg.loss.pid_alpha = .1

# center loss parameters
_cfg.loss.center_lr = 0.5
_cfg.loss.center_weight = 0.0005

# tuplet loss parameters
_cfg.loss.tuplet_s = 64
_cfg.loss.tuplet_beta = 0.1

# circle loss parameters
_cfg.loss.circle_margin = 0.25
_cfg.loss.circle_gamma = 80
_cfg.loss.circle_weight = 0.0005

_cfg.scheduler = CfgNode()
_cfg.scheduler.warmup_factor = 0.1
_cfg.scheduler.warmup_steps = 9
_cfg.scheduler.warmup_method = 'linear'
_cfg.scheduler.standup_steps = 0
_cfg.scheduler.decay_method = 'None'
_cfg.scheduler.decay_kwargs = []
_cfg.scheduler.step_by_batch = False

_cfg.test = CfgNode()
_cfg.test.feat_norm = True
_cfg.test.remove_junk = True
_cfg.test.period = 10
_cfg.test.device = 'cuda'
_cfg.test.output_dir = ""
_cfg.test.model_path = "outputs/veri776_b64_pvasff__/model_120.pth"
_cfg.test.max_rank = 50
_cfg.test.rerank = False
_cfg.test.lambda_ = 0.5
_cfg.test.debug = False
# split: When the CUDA memory is not sufficient, 
# we can split the dataset into different parts
# for the computing of distance.
_cfg.test.split = 0

_cfg.logging = CfgNode()
_cfg.logging.level = 'info'
_cfg.logging.period = 20


def get_default_configs():
    return _cfg.clone()
